import enum
import os
import re
from functools import cached_property
from typing import List, Dict, Optional

import pandas as pd
from tqdm import tqdm


def parse_column_name_string(s):
    """
    Parses a column name string to extract field_id, instance_index, and array.
    Example: 'p10001_i0_a0' -> ('10001', 0, 0)
    """
    # Initialize default values
    field_id = None
    instance_index = None
    array = None

    # Extract array (last _a followed by digits)
    array_match = re.search(r'_a(\d+)$', s)
    if array_match:
        array = int(array_match.group(1))
        s = s[:array_match.start()]  # Remove array part from string

    # Extract instance_index (last _i followed by digits)
    instance_match = re.search(r'_i(\d+)$', s)
    if instance_match:
        instance_index = int(instance_match.group(1))
        s = s[:instance_match.start()]  # Remove instance_index part from string

    # Extract field_id (everything after 'p')
    if s.startswith('p'):
        field_id = s[1:] if s[1:] else None

    return field_id, instance_index, array


class UKBioBankValueType(enum.Enum):
    """
    The value type of data-field, describing the type of variable corresponding to it.

    There are following categories:
        11-Integer - whole numbers, for example the age of a participant on a particular date;
        21-Categorical (single) - a single answer selected from a coded list or tree of mutually exclusive options, for example a yes/no choice;
        22-Categorical (multiple) - sets of answers selected from a coded list or tree of options, for instance concurrent medications;
        31-Continuous - floating-point numbers, for example the height of a participant;
        41-Text - data composed of alphanumeric characters, for example the first line of an address;
        51-Date - a calendar date, for example 14th October 2010;
        61-Time - a time, for example 13:38:05 on 14th October 2010;
        101-Compound - a set of values required as a whole to describe some compound property, for example an ECG trace;
    """
    Integer = 11
    CategoricalSingle = 21
    CategoricalMultiple = 22
    Continuous = 31
    Text = 41
    Date = 51
    Time = 61
    Compound = 101


class ValueTypeGetter:
    """
    A class to get the value type of a field from UK BioBank data.
    """

    def __init__(self, id_file_path: str):
        """
        Initializes the ValueTypeGetter.

        :param id_file_path: The path to the id.csv file.
        """
        if not os.path.exists(id_file_path):
            raise FileNotFoundError(f"The file was not found: {id_file_path}")
        self.id_file_path = id_file_path

    @cached_property
    def _fields(self) -> Dict[str, UKBioBankValueType]:
        """
        Loads the fields from the id.csv file.
        """
        fields_df = pd.read_csv(self.id_file_path)
        fields_map = {}
        for _, row in tqdm(fields_df.iterrows(), total=len(fields_df), desc="Processing fields"):
            field_id = str(row['field_id'])
            try:
                value_type = UKBioBankValueType(row['value_type'])
                fields_map[field_id] = value_type
            except (ValueError, KeyError):
                # Handle cases where value_type is not a valid UKBioBankValueType or column is missing
                if 'value_type' in row:
                    print(f"Warning: Invalid value_type '{row['value_type']}' for field_id '{field_id}'. Skipping.")
                else:
                    print(f"Warning: 'value_type' column not found for field_id '{field_id}'. Skipping.")
        return fields_map

    def get_single_value_type(self, column_name: str) -> Optional[UKBioBankValueType]:
        """
        Gets the value type for a single column name.

        :param column_name: The column name to parse and look up.
        :return: The UKBioBankValueType for the given column name, or None if not found.
        """
        field_id, _, _ = parse_column_name_string(column_name)
        if field_id:
            return self._fields.get(field_id)
        return None

    def get_batch_value_types(self, column_names: List[str]) -> List[Optional[UKBioBankValueType]]:
        """
        Gets the value types for a list of column names.

        :param column_names: A list of column names to look up.
        :return: A list of UKBioBankValueTypes corresponding to the given column names.
        """
        return [self.get_single_value_type(column_name) for column_name in column_names]


if __name__ == '__main__':
    # IMPORTANT: Please ensure the id_file_path points to your 'id.csv' file.
    # The path below is based on the provided context and may need to be adjusted.
    id_file_path = "/home/xw/python3test/omics-bert-label-encoding/datasets/raw/ukbiobank/v3/id.csv"

    value_type_getter = ValueTypeGetter(id_file_path=id_file_path)

    column_names = ["p3_i0", "pzbtb17_i2", "p23480_i1", "non_existent_field"]
    value_types = value_type_getter.get_batch_value_types(column_names)

    for column_name, value_type in zip(column_names, value_types):
        if value_type:
            print(f"Column Name: {column_name}, Value Type: {value_type.name}")
        else:
            print(f"Column Name: {column_name}, Value Type: Not Found")
